from dotenv import load_dotenv
load_dotenv()

from api import get_chatglm_response_via_sdk, get_characterglm_response
from typing import Tuple

def get_roles(content):

    prompt_template = '''
    阅读以下内容：
    """{}"""

    分析其中包含的角色和角色信息，以如下格式返回：

    角色名：
    角色信息：
    '''

    msg_content = prompt_template.format(content)

    messages = [
        {"role": "user", "content": msg_content}
    ]
    response_str = ''
    for chunk in get_chatglm_response_via_sdk(messages):
        response_str += chunk
    roles = response_str.split('\n\n')
    if (len(roles) != 2):
        print('初始内容里至少应包含两个角色')
        exit(0)

    parsed_roles = []
    for role in roles:
        parsed_roles.append(parse_role(role_info=role))
    
    return parsed_roles

def parse_role(role_info: str) -> Tuple[str, str]:
    lines = role_info.split('\n')
    role_name = lines[0].replace('角色名：', '')
    role_desc = lines[1].replace('角色信息：', '')
    print(f'role_name: {role_name}')
    print(f'role_desc: {role_desc}')
    return (role_name, role_desc)

def get_initial_chat_msg(parsed_roles, initial_content):
    prompt_template = '''
    阅读以下内容：
    """{}"""

    如果你是{}，你会对{}说什么？

    只告诉我要说的话就行，不用说其他的。
    '''

    msg_content = prompt_template.format(initial_content, parsed_roles[0][0], parsed_roles[1][0])

    messages = [
        {"role": "user", "content": msg_content}
    ]

    response_str = ''
    for chunk in get_chatglm_response_via_sdk(messages):
        response_str += chunk

    response_str = response_str.replace('"','')

    print(response_str)
    
    return response_str

def get_next_response(messages, initial_content):
    chats = []
    for msg in messages:
        chats.append(msg['content'])

    prompt_template = '''
    以下是两个人的对话内容
    """{}"""

    下一个人该如何回复？

    只告诉我要说的话就行，不用说其他的。
    '''
    messages = [
        {"role": "user", "content": prompt_template.format('\n'.join(chats))}
    ]

    response_str = ''
    for chunk in get_chatglm_response_via_sdk(messages):
        response_str += chunk

    response_str = response_str.replace('"','')

    print(response_str)
    
    return response_str

def start_chat(roles, first_msg, initial_content):
    character_meta = {
        "user_info": roles[0][-1],
        "bot_info": roles[1][-1],
        "user_name": roles[0][0],
        "bot_name": roles[1][0]
    }

    loop_count = 0
    role = 'user' if loop_count % 2 == 0 else 'assistant'

    messages = [
        {"role": role, "content": first_msg}
    ]

    for i in range(3):
        response_str = ''
        for chunk in get_characterglm_response(messages, meta=character_meta):
            response_str += chunk
        
        print(response_str)

        loop_count += 1
        role = 'user' if loop_count % 2 == 0 else 'assistant'
        messages.append(
            {"role": role, "content": response_str}
        )
        if (role == 'assistant'):
            loop_count += 1
            role = 'user' if loop_count % 2 == 0 else 'assistant'
            r = get_next_response(messages, initial_content)
            messages.append(
                {"role": role, "content": r}
            )

    return messages

def save_messages(messages, roles):
    with open('role_play.txt', 'w', encoding='utf-8') as file:
        for m in messages:
            role_index = 0 if  m['role'] == 'user' else 1
            file.write((roles[role_index][0]) + '：' + m['content'] + '\n')

    print('messages saved to role_play.txt')

    

if __name__ == "__main__":
    text = '小明，老师眼中的好学生，学习成绩优异；小强，老师眼中的坏学生，整天调皮捣蛋。小强在学校里经常欺负小明。'

    roles = get_roles(text)
    first_chat_msg = get_initial_chat_msg(roles, text)
    messages = start_chat(roles, first_chat_msg, text)
    save_messages(messages, roles)